<template>
  <div class="app-container">
    <el-container>
      <el-header>
        <h1>AI图像生成</h1>
      </el-header>
      <el-main>
        <el-row :gutter="20">
          <el-col :span="8">
            <el-scrollbar height="calc(100vh - 100px)">
              <el-card class="control-panel">
                <template #header>
                  <div class="card-header">
                    <span>生成参数</span>
                    <el-button type="primary" link @click="showAdvancedSettings = !showAdvancedSettings">
                      {{ showAdvancedSettings ? '隐藏高级设置' : '显示高级设置' }}
                    </el-button>
                  </div>
                </template>
                
                <el-form :model="form" label-position="top" size="default">
                  <!-- 基本设置 -->
                  <el-form-item label="选择模型">
                    <el-tabs v-model="activeModelTab">
                      <el-tab-pane label="基础模型" name="checkpoint">
                        <el-select 
                          v-model="currentModel" 
                          style="width: 100%"
                          :loading="switchingModel"
                          @change="handleModelChange"
                        >
                          <el-option
                            v-for="model in models.checkpoint"
                            :key="model.id"
                            :label="model.name"
                            :value="model.id"
                          >
                            <span style="float: left">{{ model.name }}</span>
                            <span style="float: right; color: #8492a6; font-size: 13px">
                              {{ model.description }}
                            </span>
                          </el-option>
                        </el-select>
                      </el-tab-pane>
                      
                      <el-tab-pane label="LoRA 模型" name="lora">
                        <el-select 
                          v-model="currentLora" 
                          style="width: 100%"
                          :loading="switchingModel"
                          @change="handleLoraChange"
                          :disabled="!currentModel"
                        >
                          <el-option
                            v-for="model in models.lora"
                            :key="model.id"
                            :label="model.name"
                            :value="model.id"
                          >
                            <span style="float: left">{{ model.name }}</span>
                            <span style="float: right; color: #8492a6; font-size: 13px">
                              {{ model.description }}
                            </span>
                          </el-option>
                        </el-select>
                        <div class="model-tip" v-if="!currentModel">请先选择基础模型</div>
                      </el-tab-pane>
                      
                      <el-tab-pane label="VAE 模型" name="vae">
                        <el-select 
                          v-model="currentVae" 
                          style="width: 100%"
                          :loading="switchingModel"
                          @change="handleVaeChange"
                          :disabled="!currentModel"
                        >
                          <el-option
                            v-for="model in compatibleVaeModels"
                            :key="model.id"
                            :label="model.name"
                            :value="model.id"
                          >
                            <span style="float: left">{{ model.name }}</span>
                            <span style="float: right; color: #8492a6; font-size: 13px">
                              {{ model.description }}
                            </span>
                          </el-option>
                        </el-select>
                        <div class="model-tip" v-if="!currentModel">请先选择基础模型</div>
                      </el-tab-pane>
                    </el-tabs>
                  </el-form-item>

                  <el-form-item label="提示词">
                    <el-input
                      v-model="form.prompt"
                      type="textarea"
                      :rows="4"
                      placeholder="请输入图像描述..."
                    />
                  </el-form-item>

                  <el-form-item label="负面提示词">
                    <el-input
                      v-model="form.negative_prompt"
                      type="textarea"
                      :rows="2"
                      placeholder="请输入不想要的元素..."
                    />
                  </el-form-item>

                  <el-form-item label="图像尺寸">
                    <el-select v-model="selectedAspectRatio" style="width: 100%" @change="handleAspectRatioChange">
                      <el-option
                        v-for="ratio in aspectRatios"
                        :key="ratio.value"
                        :label="ratio.label"
                        :value="ratio.value"
                      >
                        <span style="float: left">{{ ratio.label }}</span>
                        <span style="float: right; color: #8492a6; font-size: 13px">
                          {{ ratio.description }}
                        </span>
                      </el-option>
                    </el-select>
                  </el-form-item>

                  <!-- 高级设置 -->
                  <template v-if="showAdvancedSettings">
                    <el-divider>高级设置</el-divider>
                    
                    <el-form-item label="采样器">
                      <el-select v-model="form.scheduler" style="width: 100%">
                        <el-option-group label="推荐选项">
                          <el-option
                            v-for="scheduler in recommendedSchedulers"
                            :key="scheduler.id"
                            :label="scheduler.name"
                            :value="scheduler.id"
                          >
                            <span style="float: left">{{ scheduler.name }}</span>
                            <span style="float: right; color: #8492a6; font-size: 13px">
                              {{ scheduler.description }}
                            </span>
                          </el-option>
                        </el-option-group>
                        <el-option-group label="其他选项">
                          <el-option
                            v-for="scheduler in otherSchedulers"
                            :key="scheduler.id"
                            :label="scheduler.name"
                            :value="scheduler.id"
                          >
                            <span style="float: left">{{ scheduler.name }}</span>
                            <span style="float: right; color: #8492a6; font-size: 13px">
                              {{ scheduler.description }}
                            </span>
                          </el-option>
                        </el-option-group>
                      </el-select>
                      <div class="scheduler-tip">
                        提示：不同采样器会影响生成速度和图像质量，建议使用DPM++ 2M
                      </div>
                    </el-form-item>

                    <el-form-item label="生成数量">
                      <el-row :gutter="10">
                        <el-col :span="12">
                          <el-input-number
                            v-model="form.num_images"
                            :min="1"
                            :max="4"
                            :step="1"
                            style="width: 100%"
                          />
                        </el-col>
                        <el-col :span="12">
                          <el-tag>每次最多生成4张图片</el-tag>
                        </el-col>
                      </el-row>
                    </el-form-item>

                    <el-form-item label="随机种子">
                      <el-row :gutter="10">
                        <el-col :span="18">
                          <el-input-number
                            v-model="form.seed"
                            :min="-1"
                            :max="2147483647"
                            style="width: 100%"
                          />
                        </el-col>
                        <el-col :span="6">
                          <el-button @click="form.seed = -1">随机</el-button>
                        </el-col>
                      </el-row>
                    </el-form-item>

                    <el-form-item label="CLIP跳过层数">
                      <el-input-number
                        v-model="form.clip_skip"
                        :min="1"
                        :max="12"
                        :step="1"
                        style="width: 100%"
                      />
                    </el-form-item>

                    <el-form-item label="输出格式">
                      <el-radio-group v-model="form.image_format">
                        <el-radio label="png">PNG</el-radio>
                        <el-radio label="jpeg">JPEG</el-radio>
                      </el-radio-group>
                    </el-form-item>
                  </template>

                  <el-divider>基本参数</el-divider>

                  <el-form-item label="推理步数">
                    <el-slider
                      v-model="form.num_inference_steps"
                      :min="10"
                      :max="100"
                      :step="1"
                      show-input
                    />
                  </el-form-item>

                  <el-form-item label="AI创意程度">
                    <el-slider
                      v-model="form.guidance_scale"
                      :min="1"
                      :max="20"
                      :step="0.1"
                      show-input
                    >
                      <template #marks>
                        <template v-for="mark in guidanceMarks" :key="mark.value">
                          <div :style="{ left: mark.position + '%' }" class="guidance-mark">
                            {{ mark.label }}
                          </div>
                        </template>
                      </template>
                    </el-slider>
                    <div class="guidance-description">
                      <span class="guidance-tip">提示：数值越小AI发挥越自由，数值越大越严格遵循提示词</span>
                    </div>
                  </el-form-item>

                  <el-form-item>
                    <el-button
                      type="primary"
                      :loading="generating"
                      @click="generateImage"
                      style="width: 100%"
                    >
                      生成图像
                    </el-button>
                  </el-form-item>
                </el-form>
              </el-card>
            </el-scrollbar>
          </el-col>

          <el-col :span="16">
            <el-card class="preview-panel">
              <template #header>
                <div class="card-header">
                  <span>预览</span>
                  <span v-if="lastGeneration" class="generation-info">
                    使用模型: {{ lastGeneration.model_name }}
                  </span>
                </div>
              </template>
              
              <el-scrollbar height="calc(100vh - 160px)">
                <div class="image-preview" v-loading="generating">
                  <template v-if="generatedImages.length > 0">
                    <el-row :gutter="20">
                      <el-col 
                        v-for="(image, index) in generatedImages" 
                        :key="index"
                        :span="24 / generatedImages.length"
                      >
                        <div class="image-container">
                          <img :src="image" alt="生成的图像" />
                          <div class="image-actions">
                            <el-button type="primary" link @click="downloadImage(image)">
                              下载
                            </el-button>
                          </div>
                        </div>
                      </el-col>
                    </el-row>
                  </template>
                  <el-empty v-else description="生成的图像将显示在这里" />
                </div>
              </el-scrollbar>
            </el-card>
          </el-col>
        </el-row>
      </el-main>
    </el-container>
  </div>
</template>

<script setup>
import { ref, reactive, onMounted, computed } from 'vue'
import axios from 'axios'
import { ElMessage } from 'element-plus'
import { Refresh } from '@element-plus/icons-vue'

const API_BASE = 'http://localhost:8000'

const form = reactive({
  prompt: '',
  negative_prompt: '',
  num_inference_steps: 20,
  guidance_scale: 7.5,
  width: 512,
  height: 512,
  num_images: 1,
  scheduler: 'DPMSolverMultistep',
  seed: -1,
  clip_skip: 1,
  image_format: 'png'
})

const generating = ref(false)
const generatedImages = ref([])
const models = ref({
  checkpoint: [],
  lora: [],
  vae: []
})
const schedulers = ref([])
const currentModel = ref('')
const currentLora = ref('')
const currentVae = ref('')
const activeModelTab = ref('checkpoint')
const switchingModel = ref(false)
const lastGeneration = ref(null)
const showAdvancedSettings = ref(false)

// 预设的尺寸选项
const aspectRatios = [
  { label: '方形', value: '1:1', width: 512, height: 512, description: '适合头像和图标' },
  { label: '横向', value: '4:3', width: 640, height: 480, description: '适合风景和场景' },
  { label: '纵向', value: '3:4', width: 480, height: 640, description: '适合人像和海报' },
  { label: '宽屏', value: '16:9', width: 768, height: 432, description: '适合桌面壁纸' },
  { label: '手机', value: '9:16', width: 432, height: 768, description: '适合手机壁纸' }
]

const selectedAspectRatio = ref('1:1')

// 处理尺寸选择变化
const handleAspectRatioChange = (value) => {
  const ratio = aspectRatios.find(r => r.value === value)
  if (ratio) {
    form.width = ratio.width
    form.height = ratio.height
  }
}

// 获取可用模型列表
const fetchModels = async () => {
  try {
    const response = await axios.get(`${API_BASE}/models`)
    models.value = response.data.models
    currentModel.value = response.data.current_model
  } catch (error) {
    console.error('Error fetching models:', error)
    ElMessage.error('获取模型列表失败')
  }
}

// 获取采样器列表
const fetchSchedulers = async () => {
  try {
    const response = await axios.get(`${API_BASE}/schedulers`)
    schedulers.value = response.data.schedulers
  } catch (error) {
    console.error('Error fetching schedulers:', error)
    ElMessage.error('获取采样器列表失败')
  }
}

// 切换基础模型
const handleModelChange = async (modelId) => {
  switchingModel.value = true
  try {
    await axios.post(`${API_BASE}/models/${modelId}`)
    ElMessage.success(`已切换到模型: ${models.value.checkpoint.find(m => m.id === modelId).name}`)
    // 清除当前选择的LoRA和VAE
    currentLora.value = ''
    currentVae.value = ''
  } catch (error) {
    console.error('Error switching model:', error)
    ElMessage.error(error.response?.data?.detail || '切换模型失败')
    currentModel.value = models.value.checkpoint[0]?.id
  } finally {
    switchingModel.value = false
  }
}

// 切换LoRA模型
const handleLoraChange = async (modelId) => {
  if (!modelId) return
  switchingModel.value = true
  try {
    await axios.post(`${API_BASE}/models/${modelId}`)
    ElMessage.success(`已应用LoRA模型: ${models.value.lora.find(m => m.id === modelId).name}`)
  } catch (error) {
    console.error('Error applying LoRA:', error)
    ElMessage.error(error.response?.data?.detail || '应用LoRA模型失败')
    currentLora.value = ''
  } finally {
    switchingModel.value = false
  }
}

// 切换VAE模型
const handleVaeChange = async (modelId) => {
  if (!modelId) return
  switchingModel.value = true
  try {
    await axios.post(`${API_BASE}/models/${modelId}`)
    ElMessage.success(`已应用VAE模型: ${models.value.vae.find(m => m.id === modelId).name}`)
  } catch (error) {
    console.error('Error applying VAE:', error)
    ElMessage.error(error.response?.data?.detail || '应用VAE模型失败')
    currentVae.value = ''
  } finally {
    switchingModel.value = false
  }
}

// 获取与当前基础模型兼容的VAE模型
const compatibleVaeModels = computed(() => {
  if (!currentModel.value) return []
  return models.value.vae.filter(vae => 
    vae.compatible_models.includes(currentModel.value)
  )
})

const generateImage = async () => {
  if (!form.prompt) {
    ElMessage.warning('请输入提示词')
    return
  }

  generating.value = true
  try {
    const formData = new FormData()
    Object.keys(form).forEach(key => {
      if (form[key] !== null && form[key] !== '') {
        formData.append(key, form[key])
      }
    })

    const response = await axios.post(`${API_BASE}/generate`, formData)
    generatedImages.value = response.data.image_urls.map(url => `${API_BASE}${url}`)
    lastGeneration.value = response.data.parameters
    ElMessage.success('图像生成成功')
  } catch (error) {
    console.error('Error generating image:', error)
    ElMessage.error(error.response?.data?.detail || '生成图像时发生错误')
  } finally {
    generating.value = false
  }
}

const downloadImage = async (imageUrl) => {
  try {
    const response = await axios.get(imageUrl, { responseType: 'blob' })
    const url = window.URL.createObjectURL(new Blob([response.data]))
    const link = document.createElement('a')
    link.href = url
    link.setAttribute('download', `generated_image.${form.image_format}`)
    document.body.appendChild(link)
    link.click()
    document.body.removeChild(link)
    window.URL.revokeObjectURL(url)
  } catch (error) {
    console.error('Error downloading image:', error)
    ElMessage.error('下载图像失败')
  }
}

// 引导系数标记点
const guidanceMarks = [
  { value: 1, label: '自由发挥', position: 0 },
  { value: 7.5, label: '平衡', position: 35 },
  { value: 20, label: '严格遵循', position: 100 }
]

// 推荐的采样器
const recommendedSchedulers = computed(() => {
  return schedulers.value.filter(s => 
    ['DPMSolverMultistep', 'DPMSolverSinglestep', 'EulerA', 'DDIM'].includes(s.id)
  )
})

// 其他采样器
const otherSchedulers = computed(() => {
  return schedulers.value.filter(s => 
    !['DPMSolverMultistep', 'DPMSolverSinglestep', 'EulerA', 'DDIM'].includes(s.id)
  )
})

onMounted(() => {
  fetchModels()
  fetchSchedulers()
})
</script>

<style>
.app-container {
  min-height: 100vh;
  background-color: #f5f7fa;
}

.el-header {
  background-color: #fff;
  color: #409EFF;
  text-align: center;
  line-height: 60px;
  box-shadow: 0 2px 12px 0 rgba(0,0,0,0.1);
  position: sticky;
  top: 0;
  z-index: 100;
}

.el-main {
  padding: 20px;
  height: calc(100vh - 60px);
}

.control-panel {
  background-color: #fff;
  border-radius: 4px;
}

.preview-panel {
  height: calc(100vh - 100px);
  background-color: #fff;
  border-radius: 4px;
}

.image-preview {
  padding: 20px;
  min-height: calc(100vh - 200px);
  display: flex;
  align-items: center;
  justify-content: center;
  background-color: #fafafa;
  border-radius: 4px;
}

.image-container {
  width: 100%;
  display: flex;
  flex-direction: column;
  align-items: center;
  margin-bottom: 20px;
}

.image-container img {
  max-width: 100%;
  height: auto;
  object-fit: contain;
  border-radius: 4px;
  box-shadow: 0 2px 12px 0 rgba(0,0,0,0.1);
}

.image-actions {
  margin-top: 10px;
  display: flex;
  justify-content: center;
  gap: 10px;
}

.card-header {
  display: flex;
  justify-content: space-between;
  align-items: center;
  position: sticky;
  top: 0;
  background-color: #fff;
  z-index: 1;
}

.generation-info {
  font-size: 14px;
  color: #909399;
}

.el-select-dropdown__item {
  height: auto;
  padding: 8px;
}

.el-divider {
  margin: 16px 0;
}

.el-form-item {
  margin-bottom: 18px;
}

.el-tag {
  margin-left: 10px;
}

.aspect-ratio-button {
  width: 100%;
  margin-bottom: 10px;
}

:deep(.el-form-item__label) {
  font-weight: 500;
}

:deep(.el-card__body) {
  padding: 20px;
}

.el-scrollbar {
  --el-scrollbar-opacity: 0.3;
  --el-scrollbar-hover-opacity: 0.5;
}

.el-scrollbar__bar {
  z-index: 2;
}

.mb-2 {
  margin-bottom: 8px;
}

.dimension-input {
  display: flex;
  flex-direction: column;
  gap: 4px;
}

.dimension-label {
  font-size: 12px;
  color: #606266;
}

.dimension-swap-btn {
  width: 100%;
  padding: 12px;
}

:deep(.el-input-number .el-input__wrapper) {
  padding-right: 0;
}

.guidance-description {
  margin-top: 8px;
  font-size: 12px;
  color: #909399;
  text-align: center;
}

.guidance-tip {
  color: #606266;
}

.guidance-mark {
  position: absolute;
  transform: translateX(-50%);
  font-size: 12px;
  color: #909399;
  margin-top: 4px;
}

.model-tip {
  margin-top: 8px;
  font-size: 12px;
  color: #909399;
  text-align: center;
}

:deep(.el-tabs__content) {
  padding: 16px 0;
}

.scheduler-tip {
  margin-top: 8px;
  font-size: 12px;
  color: #909399;
  text-align: center;
}

:deep(.el-select-group__title) {
  padding: 8px;
  font-size: 13px;
  font-weight: bold;
  color: #606266;
}
</style> 